from mitmproxy import contentviews
from mitmproxy import flow, ctx, tls
from mitmproxy import http
from mitmproxy.addonmanager import Loader
from mitmproxy import connection
import mitmproxy.proxy.server_hooks

import typing, json

from urllib.parse import urljoin, urlparse

class MitmAddon:
    def __init__(self):
        self.snis: typing.Set[str] = set()
        self.failed_snis: typing.Set[str] = set()
        self.failed_ips: typing.Set[str] = set()
        self.connections: typing.List[typing.Dict] = []
        self.flows: typing.List[typing.Dict] = []

    def load(self, loader: Loader):
        loader.add_option(
            "output_json_path", typing.Optional[str], None, "Output JSON file path"
        )

    async def running(self):
        ctx.log.info("AndroidMitmAddon is running")

    async def done(self):
        ctx.log.info("AndroidMitmAddon is shutting down")
        if ctx.options.output_json_path:
            with open(ctx.options.output_json_path, "w") as f:
                json.dump({
                    'snis': list(self.snis),
                    'failed_snis': list(self.failed_snis),
                    'failed_ips': list(self.failed_ips),
                    'connections': self.connections,
                    'flows': self.flows,
                }, f)
            ctx.log.info("Output written to {}".format(ctx.options.output_json_path))

    def tls_failed_client(self, data: tls.TlsData):
        message = data.conn.error
        if message is not None and "certificate" not in message:
            return

        sni = data.conn.sni
        if sni is not None and sni not in self.failed_snis:
            ctx.log.info(
                "Client-side handshake failed for {}, marking it as ignored. Error message: {}".format(
                    sni, message
                )
            )
            self.failed_snis.add(sni)
            return

        remote = data.context.server.address
        if remote is not None:
            ctx.log.info(
                "Client-side handshake failed for {}, marking it as ignored. Error message: {}".format(
                    remote[0], message
                )
            )
            self.failed_ips.add(remote[0])

    def tls_clienthello(self, data: tls.ClientHelloData):
        data.establish_server_tls_first = True

        sni = data.client_hello.sni
        if sni is not None:
            self.snis.add(sni)
            if sni in self.failed_snis:
                ctx.log.info("Ignoring connection to {} as it does not trust our certificate".format(sni))
                data.ignore_connection = True
            return

        remote = data.context.server.address
        if remote is not None and remote[0] in self.failed_ips:
            ctx.log.info("Ignoring connection to {} as it does not trust our certificate".format(remote[0]))
            data.ignore_connection = True

    def server_connected(self, data: mitmproxy.proxy.server_hooks.ServerConnectionHookData):
        self._record_connection(data.server, "connected")

    def server_connect_error(self, data: mitmproxy.proxy.server_hooks.ServerConnectionHookData):
        self._record_connection(data.server, "failed")

    def _record_connection(self, conn: connection.Server, status: str):
        c = {
            'status': status,
        }
        if conn.peername is not None:
            c['ip'] = conn.peername[0]
            c['port'] = conn.peername[1]
        if conn.address is not None:
            c['host'] = conn.address[0]
        if conn.error is not None:
            c['error'] = conn.error
        self.connections.append(c)

    def response(self, flow: http.HTTPFlow):
        self._record_flow(flow)

    def error(self, flow: http.HTTPFlow):
        self._record_flow(flow)

    def _record_flow(self, flow: http.HTTPFlow):
        f = {
            'request': {
                'host': flow.request.pretty_host,
                'method': flow.request.method,
                'path': flow.request.path,
                'scheme': flow.request.scheme,
            }
        }
        if flow.server_conn.timestamp_start is not None:
            # We have a server connection
            f['connection'] = {
                'ip': flow.server_conn.peername[0],
                'port': flow.server_conn.peername[1],
            }
        if flow.response is not None:
            # We have a response
            f['response'] = {
                'status_code': flow.response.status_code,
                'reason': flow.response.reason,
            }
        if flow.websocket is not None:
            # Mark this flow as having a WebSocket connection
            f['websocket'] = True
        if flow.error is not None:
            # We have an error
            f['error'] = flow.error.msg
        self.flows.append(f)

addons = [MitmAddon()]